package com.web;

import lombok.SneakyThrows;

import javax.crypto.Cipher;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * @author：zhang guoxiang
 * @date：2022/5/5
 */
public abstract class RSAUtils {

    /**
     * RSA最大加密明文大小
     */
    private static final int MAX_ENCRYPT_BLOCK = 117;

    /**
     * RSA最大解密密文大小
     */
    private static final int MAX_DECRYPT_BLOCK = 128;

    /**
     * 获取密钥对
     *
     * @return 密钥对
     */
    @SneakyThrows
    public static KeyPair getKeyPair() {
        KeyPairGenerator generator = KeyPairGenerator.getInstance("RSA");
        generator.initialize(2048);
        return generator.generateKeyPair();
    }

    /**
     * 获取私钥
     *
     * @param privateKey 私钥字符串
     * @return
     */
    @SneakyThrows
    public static PrivateKey getPrivateKey(String privateKey) {
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        byte[] decodedKey = Base64.getDecoder().decode(privateKey.getBytes());
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(decodedKey);
        return keyFactory.generatePrivate(keySpec);
    }

    /**
     * 获取公钥
     *
     * @param publicKey 公钥字符串
     * @return
     */
    @SneakyThrows
    public static PublicKey getPublicKey(String publicKey) {
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        byte[] decodedKey = Base64.getDecoder().decode(publicKey.getBytes());
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(decodedKey);
        return keyFactory.generatePublic(keySpec);
    }

    /**
     * RSA加密
     *
     * @param data      待加密数据
     * @param publicKey 公钥
     * @return
     */
    @SneakyThrows
    public static String encrypt(String data, PublicKey publicKey) {
        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        int inputLen = data.getBytes().length;
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        int offset = 0;
        byte[] cache;
        int i = 0;
        // 对数据分段加密
        while (inputLen - offset > 0) {
            if (inputLen - offset > MAX_ENCRYPT_BLOCK) {
                cache = cipher.doFinal(data.getBytes(), offset, MAX_ENCRYPT_BLOCK);
            } else {
                cache = cipher.doFinal(data.getBytes(), offset, inputLen - offset);
            }
            out.write(cache, 0, cache.length);
            i++;
            offset = i * MAX_ENCRYPT_BLOCK;
        }
        byte[] encryptedData = out.toByteArray();
        out.close();
        // 获取加密内容使用base64进行编码,并以UTF-8为标准转化成字符串
        // 加密后的字符串
        return Base64.getEncoder().encodeToString(encryptedData);
    }

    /**
     * RSA解密
     *
     * @param data       待解密数据
     * @param privateKey 私钥
     * @return
     */
    @SneakyThrows
    public static String decrypt(String data, PrivateKey privateKey) {
        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        byte[] dataBytes = Base64.getDecoder().decode(data);
        int inputLen = dataBytes.length;
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        int offset = 0;
        byte[] cache;
        int i = 0;
        // 对数据分段解密
        while (inputLen - offset > 0) {
            if (inputLen - offset > MAX_DECRYPT_BLOCK) {
                cache = cipher.doFinal(dataBytes, offset, MAX_DECRYPT_BLOCK);
            } else {
                cache = cipher.doFinal(dataBytes, offset, inputLen - offset);
            }
            out.write(cache, 0, cache.length);
            i++;
            offset = i * MAX_DECRYPT_BLOCK;
        }
        byte[] decryptedData = out.toByteArray();
        out.close();
        // 解密后的内容
        return new String(decryptedData, StandardCharsets.UTF_8);
    }


    /**
     * 签名
     *
     * @param data       待签名数据
     * @param privateKey 私钥
     * @return 签名
     */
    @SneakyThrows
    public static String sign(String data, PrivateKey privateKey) {
        byte[] keyBytes = privateKey.getEncoded();
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        PrivateKey key = keyFactory.generatePrivate(keySpec);
        Signature signature = Signature.getInstance("MD5withRSA");
        signature.initSign(key);
        signature.update(data.getBytes());
        return Base64.getEncoder().encodeToString(signature.sign());
    }

    /**
     * 验签
     *
     * @param srcData   原始字符串
     * @param publicKey 公钥
     * @param sign      签名
     * @return 是否验签通过
     */
    @SneakyThrows
    public static boolean verify(String srcData, PublicKey publicKey, String sign) {
        byte[] keyBytes = publicKey.getEncoded();
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        PublicKey key = keyFactory.generatePublic(keySpec);
        Signature signature = Signature.getInstance("MD5withRSA");
        signature.initVerify(key);
        signature.update(srcData.getBytes());
        return signature.verify(Base64.getDecoder().decode(sign.getBytes()));
    }

    /**
     * 文件解密
     *
     * @param publicKeyStr
     * @param srcFileName
     * @param destFileName
     */
    @SneakyThrows
    public static void decryptFileBig(String publicKeyStr, String srcFileName, String destFileName) {
        Cipher cipher = null;
        try (InputStream is = new FileInputStream(srcFileName); OutputStream out = new FileOutputStream(destFileName)) {
            X509EncodedKeySpec encodedKey = new X509EncodedKeySpec(Base64.getDecoder().decode(publicKeyStr));
            KeyFactory keyf = KeyFactory.getInstance("RSA");
            PublicKey pubKey = keyf.generatePublic(encodedKey);
            cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.DECRYPT_MODE, pubKey);
            File f = new File(srcFileName);
            int blockSize = cipher.getBlockSize();
            int size = Integer.parseInt(String.valueOf(f.length()));
            byte[] decryptByte = new byte[size];
            is.read(decryptByte);
            //分别对各块数据进行解密
            int j = 0;
            while ((decryptByte.length - j * blockSize) > 0) {
                out.write(cipher.doFinal(decryptByte, j * blockSize, blockSize));
                j++;
            }
        }
    }

    /**
     * 加密大文件
     *
     * @param privateKeyStr
     * @param srcFileName
     * @param destFileName
     */
    @SneakyThrows
    public static void encryptFileBig(String privateKeyStr, String srcFileName, String destFileName) {
        if (privateKeyStr == null) {
            throw new RuntimeException("加密私钥为空, 请设置");
        }
        Cipher cipher = null;
        try (InputStream is = new FileInputStream(srcFileName); OutputStream out = new FileOutputStream(destFileName)) {
            // 使用默认RSA
            PKCS8EncodedKeySpec priPKCS8 = new PKCS8EncodedKeySpec(
                    Base64.getDecoder().decode(privateKeyStr));
            KeyFactory keyf = KeyFactory.getInstance("RSA");
            PrivateKey priKey = keyf.generatePrivate(priPKCS8);
            cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.ENCRYPT_MODE, priKey);
            int blockSize = cipher.getBlockSize();
            File f = new File(srcFileName);
            int size = Integer.parseInt(String.valueOf(f.length()));
            byte[] encryptByte = new byte[size];
            is.read(encryptByte);
            int outputBlockSize = cipher.getOutputSize(encryptByte.length);
            int leavedSize = encryptByte.length % blockSize;
            int blocksNum = leavedSize == 0 ? encryptByte.length / blockSize : encryptByte.length / blockSize + 1;
            byte[] cipherData = new byte[blocksNum * outputBlockSize];
            for (int i = 0; i < blocksNum; i++) {
                cipher.doFinal(encryptByte, i * blockSize, Math.min((encryptByte.length - i * blockSize), blockSize), cipherData, i * outputBlockSize);
            }
            out.write(cipherData);
        }
    }
}
